"""
This script implements A* path planning and visualization for navigating a grid-based environment. 
The main features include:

1. **Node Class**:
   - Represents individual nodes in the A* algorithm with cost calculations (g, h, f).
2. **A* Algorithm**:
   - Finds the shortest path between a start and goal position in a grid map.
3. **Path Planning Through Multiple Goals**:
   - Plans a path sequentially through a list of goals using A*.
4. **Path Visualization**:
   - Draws the computed path on a grid map using Matplotlib.

Utility functions are provided to convert between world and grid coordinates. This script is modular 
and can be integrated into larger systems for navigation and visualization tasks.
"""

import numpy as np
import heapq
import matplotlib.pyplot as plt

class Node:
    """
    A class representing a node in the A* algorithm.

    Attributes:
    position -- Tuple (x, y) representing the position of the node.
    parent -- Parent node for path reconstruction (default: None).
    g -- Cost from the start node to the current node.
    h -- Heuristic cost from the current node to the goal node.
    f -- Total cost (f = g + h).

    Methods:
    __lt__ -- Defines a less-than operator for priority queue comparison.
    """
    def __init__(self, position, parent=None):
        self.position = position
        self.parent = parent
        self.g = 0
        self.h = 0
        self.f = 0

    def __lt__(self, other):
        return self.f < other.f


def a_star(grid_map, start, goal, grid_size):
    """
    Perform path planning using the A* algorithm.

    Parameters:
    grid_map -- 2D numpy array representing the grid map (1: obstacle, 0: free space).
    start -- Tuple (x, y) representing the start position.
    goal -- Tuple (x, y) representing the goal position.
    grid_size -- Size of each grid cell.

    Returns:
    List of nodes representing the planned path, or None if no path is found.
    """
    # Initialize open and closed lists
    open_list = []
    open_dict = {}
    closed_list = set()

    # Add the start node to the open list
    start_node = Node(start)
    goal_node = Node(goal)
    heapq.heappush(open_list, start_node)
    open_dict[start] = start_node

    # Define movement directions (up, down, left, right)
    directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]

    while open_list:
        # Pop the node with the smallest f value
        current_node = heapq.heappop(open_list)
        open_dict.pop(current_node.position, None)

        # Add the current node to the closed list
        closed_list.add(tuple(current_node.position))

        # Check if the goal is reached
        if np.array_equal(current_node.position, goal_node.position):
            path = []
            while current_node:
                path.append((current_node.position[1], current_node.position[0]))  # Convert (row, column) to (x, y)
                current_node = current_node.parent
            return path[::-1]  # Return the reversed path

        # Process neighbors
        for direction in directions:
            neighbor_pos = (current_node.position[0] + direction[0], current_node.position[1] + direction[1])

            # Check for out-of-bounds or obstacles
            if (neighbor_pos[0] < 0 or neighbor_pos[0] >= grid_map.shape[0] or
                neighbor_pos[1] < 0 or neighbor_pos[1] >= grid_map.shape[1] or
                grid_map[neighbor_pos] == 1):
                continue

            # Skip if the neighbor is already in the closed list
            if tuple(neighbor_pos) in closed_list:
                continue

            # Compute g, h, f values
            g = current_node.g + 1
            h = abs(neighbor_pos[0] - goal_node.position[0]) + abs(neighbor_pos[1] - goal_node.position[1])
            f = g + h

            # Check and update the node in the open list
            if neighbor_pos in open_dict:
                existing_node = open_dict[neighbor_pos]
                if g < existing_node.g:
                    existing_node.g = g
                    existing_node.h = h
                    existing_node.f = f
                    existing_node.parent = current_node
                    heapq.heapify(open_list)
            else:
                neighbor_node = Node(neighbor_pos, current_node)
                neighbor_node.g = g
                neighbor_node.h = h
                neighbor_node.f = f
                heapq.heappush(open_list, neighbor_node)
                open_dict[neighbor_pos] = neighbor_node

    return None  # No path found


def plan_path_through_goals(pos0, goals, freespace_map, min_x, max_x, min_y, max_y, grid_size, ax, visualize=False):
    """
    Plan a path through a series of goals using A*.

    Parameters:
    pos0 -- Starting position [x, y, theta].
    goals -- List of goal positions [x, y, theta].
    freespace_map -- 2D numpy array representing the free space map.
    min_x, max_x, min_y, max_y -- Boundaries of the map in the world coordinate system.
    grid_size -- Size of each grid cell.
    ax -- Matplotlib Axes object for visualization.
    visualize -- Boolean flag to enable visualization.

    Returns:
    List of path points in world coordinates.
    """
    full_path = []
    current_pos = pos0

    for goal in goals:
        path_grid = a_star(freespace_map,
                           world_to_grid(current_pos[:2], min_x, max_x, min_y, max_y, grid_size)[::-1],
                           world_to_grid(goal[:2], min_x, max_x, min_y, max_y, grid_size)[::-1],
                           grid_size)

        if path_grid:
            path_astar = [grid_to_world(point, min_x, min_y, grid_size) for point in path_grid]
            full_path.extend(path_astar)

            if visualize:
                draw_path(ax, path_astar, grid_size, min_x, min_y)
            
            current_pos = goal
        else:
            print(f"No A* path found to goal {goal}.")
            break

    return full_path


def draw_path(ax, path, grid_size, min_x, min_y):
    """
    Draw a path on an existing plot.

    Parameters:
    ax -- Matplotlib Axes object.
    path -- List of path points in world coordinates.
    grid_size -- Size of each grid cell.
    min_x, min_y -- Bottom-left corner of the map in world coordinates.

    Returns:
    None.
    """
    path_x = [p[0] for p in path]
    path_y = [p[1] for p in path]

    ax.plot(path_x, path_y, color="purple", marker="o", markersize=0.5, label="A* Path")
    plt.draw()


def grid_to_world(grid_pos, min_x, min_y, grid_size):
    """
    Convert grid coordinates to world coordinates.

    Parameters:
    grid_pos -- Tuple (grid_x, grid_y).
    min_x -- Minimum x-coordinate of the map.
    min_y -- Minimum y-coordinate of the map.
    grid_size -- Size of each grid cell.

    Returns:
    Numpy array [x, y].
    """
    x = min_x + grid_pos[0] * grid_size
    y = min_y + grid_pos[1] * grid_size
    return np.array([x, y])


def world_to_grid(world_pos, min_x, max_x, min_y, max_y, grid_size):
    """
    Convert world coordinates to grid coordinates.

    Parameters:
    world_pos -- Tuple (x, y).
    min_x, max_x, min_y, max_y -- Boundaries of the map in world coordinates.
    grid_size -- Size of each grid cell.

    Returns:
    Tuple (grid_x, grid_y).
    """
    grid_x = int((world_pos[0] - min_x) / grid_size)
    grid_y = int((world_pos[1] - min_y) / grid_size)

    grid_x = max(0, min(grid_x, int((max_x - min_x) / grid_size)))
    grid_y = max(0, min(grid_y, int((max_y - min_y) / grid_size)))

    return grid_x, grid_y
